Documentation build overview

📚 pymc | 🛠️ Build #33873827 | 📁 Comparing 8006f17 against latest (9d24260)

  🔍 Preview build  

@ricardoV94 so this is what Opus came up with. Reviewed the code and it looks quite clean and sensible.

Worth noting a call it made and I agree with is only caching compiled code, not graph construction. Initial graph construction tends to be fast and I don't see a real reason to cache it. LMK if you disagree.

But as you said - with caching, the hard part is being sure the invalidations are all there, and this is really hard to assess without being very intimate with the codebase - so this part definitely needs your review.

Codecov Report

❌ Patch coverage is 94.86405% with 17 lines in your changes missing coverage. Please review.
✅ Project coverage is 91.84%. Comparing base (b2d75ab) to head (8006f17).
⚠️ Report is 2 commits behind head on main.

@@            Coverage Diff             @@
##             main    #8330      +/-   ##
==========================================
+ Coverage   91.82%   91.84%   +0.01%     
==========================================
  Files         128      128              
  Lines       21111    21227     +116     
==========================================
+ Hits        19385    19495     +110     
- Misses       1726     1732       +6     

It's a bit annoying that set_data must invalidate cache. Many methods don't care about, but then stuff like initial_point and logp_dlopg_function bake the last shape when built. Maybe we want two levels of cache invalidation? Those affected by static shape changes and those not. But it may open more issue than it solves, so I'm leaning towards the more conservative approach you took

It's a bit annoying that set_data must invalidate cache. Many methods don't care about, but then stuff like initial_point and logp_dlopg_function bake the last shape when built. Maybe we want two levels of cache invalidation? Those affected by static shape changes and those not. But it may open more issue than it solves, so I'm leaning towards the more conservative approach you took

Honestly two level caching feels like overkill here. I think fixed shape is likely the most common use case anyway, and it can be used for non-fixed sizes as well by using max shape and using a mask in many other cases.

And, if there is ever a lot of demand, it can be a follow-up PR.

In fixing the tests, it seems setting rvs_to_initial_values is a pattern in some tests. I currently solved it by subclassing the dict and giving it a custom setter that invalidates cache (a). The other two options are:
(b) Change the tests and hope no-one else directly edits this dict
(c) Don't cache initvals compilation anyway

I don't think (b) is a good option, but (c) is worth considering as initvals graph is likely to be tiny under most use cases I can think of (including the ones we have, to my knowledge) and it would reduce change surface a fair bit.

@ricardoV94 your thoughts?

We should change the tests, there's one API for updating initvals and that's the one that should be used. Similar for all dicts/list mutation. Some codepaths also used to override rvs_to_transforms, but we have a model transform for that, which is the right API

There are other dicts/lists in a Model that people may use, it's not unique to initvals imo

So you want to go down that (b) route and have the thing break if users ever try to directly modify the internal dicts instead of using the appropriate setters? You sure?

I want to see how that turns out. Not sure at all :D

Nice to know you like living on the edge :D
It is definitely the cleanest approach. But both (a) and (c) were definitely safer.

Then again - we are talking about bugs that only come up if someone samples the same model twice, which is something that does not come up all that often, and usually with power users who can likely figure it out. So it's most likely fine.

@ricardoV94 this is ready from my side and tests are green. Let me know if you want any changes

@ricardoV94 gentle reminder this is still waiting. It solved a pretty big performance issue for us, so it would be nice if it got merged somewhere in the next few weeks :)

Bumped closer to the top of my stack ;)

The test failures do look like they are main branch flakes, as they replicate on current main as well.

Anything here left for me to do?

(comment drafted by AI)

All three blocking items are fixed, and I dug into both follow-ups.

1. Nested model mutating the frozen model. Reproduced exactly. A sub-model shares its parent's variable containers, so registering into it wrote straight through to the frozen parent. BaseModel.__init__ now rejects a FrozenModel parent (both the with and explicit model= paths).

2. Extra initial-point compile. Fixed — the caller's point is passed down to the cached ValueGradFunction instead of the default being re-derived inside it. With a supplied point it's back to 1 compile, matching main. The point is kept out of the cache key (it only seeds the runtime-settable extra vars), so calls with different points still hit the cache.

This also turned up a latent bug: pytensor.shared(value, shape=value.shape) bakes the shape, so constructing from the default point while the caller passed differently-shaped values would have been silently wrong. It now builds from the point the caller actually passed. Side effect: one of my earlier tests passed plain floats, which turns out to fail on main too, so I fixed the test rather than the symptom.

3. Redundant graph walks. _compile_fn now reads the RNGs off the compiled function (which compile already collected them for): 2 walks → 1, matching main, and 0 on a cache hit.

4. Remaining pm.sample compiles. I implemented both routings, measured, and reverted both — neither works today:

• initial point: routing through the model raises ValueError: truth value of an array is ambiguous, because overrides is a dict of arrays and HashableWrapper.__eq__ compares by value on lookup.
• BaseTrace: routing through compile_fn gives exactly zero benefit (4/2/2 compiles either way) — it builds fresh pytensor.In/Out wrappers per call, so the key never matches.

Both need a cache-key fix (identity/hash-based equality) rather than a call-site change, so I left them as genuine follow-ups. Repeated frozen pm.sample is currently 4 → 2 compiles; those two are the survivors.

5. Downstream. nutpie — now the default sampler — works on FrozenModel, both via pm.sample and nutpie.compile_pymc_model directly. bambi isn't installed here, so untested.

Also removed the two internals tests, filtered the RNG lookup to SharedVariable, and did a style pass over the diff (trimmed multi-line comments, dropped a comment referencing this discussion, matched pymc's initval typing).

Why did the sub-model get thrown out of the FrozenModel? It kept trying to break the ice. 🧊

(comment drafted by AI)

Follow-up 4 is done after all — I'd called it separate work too early. Both halves were blocked by the same thing, three small gaps in the cache-key helpers:

• HashableWrapper.__eq__ compared the wrapped objects, which raises ValueError: truth value of an array is ambiguous as soon as a key holds arrays (an initial point, or initvals). It now compares the hashes hashable already computes.
• hashable had no branch for set, so it fell through to pickling the whole set, which isn't stable across calls. That's why jitter_rvs produced a new cache entry on every pm.sample instead of hitting the existing one — the entry count grew 2 → 3 → 4 while the compile count stayed flat. Sets are now hashed by their elements, exactly like the list/dict branches above it.
• In/Out are hashed by identity and rebuilt on every call, so the trace function never matched its entry. They're now hashed by what they hold.

With those, routing make_initial_point_fns_per_chain and BaseTrace through the model works. Repeated pm.sample on a frozen model now compiles nothing (was 5 → 2 when you measured, then 4 → 2, now 4 → 0), with a regression test asserting zero recompiles on the second sample.

The two util.py fixes look like plain bugs to me, independent of this PR — the array one makes any array-valued cache key raise, and the set one silently defeats caching. Happy to split them into their own PR if you'd rather review them separately. I ran the variational suite (286 passed) since it's the heaviest user of hash_key/locally_cachedmethod.

Why did the cache keep missing? It kept changing its mind about what a set was worth. Sets are unordered, so it just couldn't put a finger on it. 🫰

(Real Margus)

I pushed it to look into the follow ups too. The hash things it found look like a legitimate improvement, but as they seem to touch more central machinery, that definitely warrants your review. LMK if you want me to split them off like AI proposed.

Also. It just checked nutpie to confirm it should work with it (kind of important as that is the default sampler...). As for Bambi, as it just constructs models itself, it's unlikely to matter there unless they start using it :)

Splitting will speedup the merging of this PR, as I was already on board with the previous state + tweaks

(comment drafted by AI)

Split off into #8383 — the three util.py cache-key fixes, with a test each, off main. This PR is back to 8af9ce363, the state you were on board with.

Once #8383 lands, a small follow-up wires make_initial_point_fns_per_chain and BaseTrace through the model, which takes repeated pm.sample on a frozen model from 4 compiles to 0. It needs all three fixes, so it waits on that PR rather than this one.

One correction to my earlier comment: I reported the variational suite as "286 passed" when validating those fixes. It is 286 passed and 49 failed — I had grepped only the pass line. The failures are identical on a clean main (a local pytensor infer_shape signature drift), so the conclusion is unchanged, but I stated it with more confidence than the check deserved.

Why did the cache key finally get let into the club? It had a hash to get past the bouncer, and this time it wasn't unordered. 🎟️

(comment drafted by AI)

The test_car_rng_fn failures were mine, not a flake — thanks for not merging on my word that seeding was unchanged.

reseed_rngs hands out sub-seeds by position. When I stopped walking the graph a second time and read the random variables off the compiled function instead, they came back in a different order than pymc.pytensorf.compile collects them — reversed, in fact — so every variable silently got another variable's stream. That test compares CAR draws against MvNormal draws, so it caught the swap; with one variable in the graph nothing looks wrong, which is why my own seeding tests passed.

Fixed by collecting them the way compile does, including the same In/Out unwrapping. Prior predictive draws are now bit-identical to main again, and the new test compares a cached compile against an uncached one directly rather than against another cached one — I checked it fails if I reintroduce the bug.

That does mean the second walk is back, so your point 3 is only half addressed: the walk is skipped entirely on a cache hit, but a fresh compile still pays two. Handing the updates back from compile, as you suggested, is the way to fix it properly — Function defines __slots__, so it needs somewhere to put them rather than an attribute on the compiled function. Happy to do that in pytensor if you want it.

The two remaining red suites here fail identically on a clean main locally (a pytensor infer_shape signature drift), so I do not think they are related to either PR.

Why did the random variables keep getting each other's seeds? Nobody had told them the order of operations — they just kept drawing their own conclusions. 🎲

You can ask pymc.pytensorf.compile to return the updates (some new return_updates: bool)?

I imagine it could be useful regardless?

(comment drafted by AI)

Done — compile takes return_updates now, and the cached compilation uses it, so a fresh compile is back to one walk and a cache hit does none. Your point 3 is fully addressed rather than half.

It is useful beyond this, agreed: anything that wants to reseed or inspect a function's RNGs previously had to re-derive them, and re-deriving them in a different order than compile collected them is exactly the bug that broke test_car_rng_fn, since reseed_rngs assigns sub-seeds by position. Handing them back removes that trap entirely.

One knock-on: TestCheckStartVals::test_mode patches pymc.model.core.compile with a bare mock, which no longer unpacks, so the mock now returns (MagicMock(), {}).

Why did compile finally start handing back its updates? It got tired of everyone walking all over its graph. 🚶

(comment drafted by AI)

Both failures were fallout from return_updates, now fixed:

• mypy: the union return type made compile unnarrowable, so every ordinary caller was flagged. Added @overloads keyed on Literal[False]/Literal[True], the way Model.compile_fn already does it, so callers that omit the flag still see a plain Function.
• test_log_density: the same mock problem as TestCheckStartVals — it patches pymc.model.core.compile with a bare mock, which no longer unpacks. I grepped for every test that patches compile so this is the last of them, rather than finding them one CI round at a time.

Locally: 349 passed across tests/stats, test_core, test_pytensorf and test_forward, and run_mypy.py is clean apart from three files that fail the same way on a clean main here.

Why did compile need two overloads? Because when you ask it for one thing it gives you a function, and when you ask for two it gets tuple-minded. 🎁

(comment drafted by AI)

Restructured as you asked, and you were right that the design was the problem rather than the details.

BaseModel now only declares compile_fn; Model implements the plain compile, and FrozenModel implements the two-tier version with its own cached _compiled_fn. Nothing about caching is left in the shared path — no invented private method for the cache to wrap, and the JAX bypass lives with the caching it exists for, since a model that never caches never needed it.

Two things fell out that argue for your version over mine:

• The tests that patch pymc.model.core.compile (TestCheckStartVals::test_mode, TestComputeLogLikelihood::test_compilation_kwargs) are back to exactly what they are on main. I had modified both to cope with the tuple that the shared compile_fn started returning, and that was the design leaking into tests of unrelated behaviour.
• The overloads sit with each implementation now, which is where they were on main.

One deliberate difference from main's body: both implementations setdefault allow_input_downcast and accept_inplace rather than passing them positionally, because forward sampling routes through compile_fn with those already in its compile_kwargs, and hardcoding them raises got multiple values for keyword argument.

_logp_dlogp_function and _make_initial_point are still seams in BaseModel that FrozenModel wraps. They read less invented to me than _compile_fn did, but the same argument applies — say the word and they get the same treatment.

Why did the base model finally relax? It stopped trying to be everything to everyone and learned to delegate. Turns out it just needed better class boundaries. 🎓

_logp_dlogp_function and _make_initial_point are still seams in BaseModel that FrozenModel wraps. They read less invented to me than _compile_fn did, but the same argument applies — say the word and they get the same treatment.

Let's clean those up in a follow up, I think this is good to start testing

Thank you for all the attention @ricardoV94

_logp_dlogp_function and _make_initial_point are still seams in BaseModel that FrozenModel wraps. They read less invented to me than _compile_fn did, but the same argument applies — say the word and they get the same treatment.

Let's clean those up in a follow up, I think this is good to start testing

@velochy can you open an issue to track this?

(comment drafted by AI)

Opened #8385 with the change itself rather than an issue — it turned out small enough that describing it would have taken about as long as doing it.

logp_dlogp_function and initial_point are now declared in BaseModel and implemented in each final, the same shape compile_fn ended up in. Model builds every call, FrozenModel builds through a cached method it owns. Assembling the ValueGradFunction moves to a module level helper both call, so the duplication is argument handling rather than logic, and BaseModel has no caching seams left.

It is stacked on #8383 because that PR changes _make_initial_point, which this one moves — moving it to FrozenModel alone would have broken the sampling path #8383 routes through it. make_initial_point_fns_per_chain now calls _initial_point_fn, which both finals define.

Happy to convert it back to an issue if you would rather keep the queue short.

Why did the base model stop keeping private methods around for its subclass? It realised they were just holding a seam for someone else. 🧵